#include <linuxboot/types.h>
#include <linuxboot/libc.h>
#include <lk/api.h>
#include "bootimg.h"

#include "u3phy.h"

#include <drivers.h>

#define INFRASYS_BASE 0x10001000

static void (*set_active_uart)(void) = (void*)(0x46002DD0|1);
static uint32_t (*get_active_uart)(void) = (void*)(0x46002BEC|1);
static void (*test)(void) = (void*)(0x46002DE0|1);

uint32_t g_boot, g_recovery;
int (*read_org)(part_dev_t *dev, uint64_t src, void *dst, int size, unsigned int part_id);

int read(part_dev_t *dev, uint64_t block_off, void *dst, int sz, unsigned int part_id)
{
    dprintf("read_hook 0x%08X from 0x%08X\n", read_org, __builtin_return_address(0));
    video_printf("read_hook 0x%08X from 0x%08X\n", read_org, __builtin_return_address(0));

    int ret = 0;
    if (block_off == g_boot * 0x200 || block_off == g_recovery * 0x200) {
        video_printf("demangle boot image - from 0x%08X\n", __builtin_return_address(0));

        if (sz < 0x400) {
            ret = read_org(dev, block_off + 0x400, dst, sz, part_id);
        } else {
            void *second_copy = (char*)dst + 0x400;
            ret = read_org(dev, block_off, dst, sz, part_id);
            memcpy(dst, second_copy, 0x400);
            memset(second_copy, 0, 0x400);
        }
    } else {
        ret = read_org(dev, block_off, dst, sz, part_id);
    }

    return ret;
}

void trap(uint32_t addr)
{
    if (addr % 4 == 0)
        ((uint32_t *)addr)[0] = 0xFFFFFFFF;
    else if (addr % 2 == 0)
        ((uint16_t *)addr)[0] = 0xFFFF;
    else
    {
        video_printf("trap: addr not divisible by 4 nor 2");
        for (;;)
            ;
    }
}

void secure_boot_disable() {
    /**((uint16_t*)0x46026B1E) = 0xBF00;
    *((uint16_t*)0x46026B28) = 0xBF00;
    *((uint16_t*)0x46026AFA) = 0xBF00;
    *((uint16_t*)0x46026B04) = 0xBF00;
    *((uint16_t*)0x46026B0E) = 0xBF00;
    *((uint16_t*)0x46026B1E) = 0xBF00;*/

    *((uint16_t*)0x4605B8EC) = 0xF04F;
    *((uint16_t*)0x4605B8EC+2) = 0x0;
    *((uint16_t*)0x4605B8EC+4) = 0x4770;
}

#if 1
#define ACR4 0x11290810
#define ACR6 0x11290818
#define DTM0 0x11290868
#define DTM1 0x1129086C

char is_usb_uart_enabled() {
    uint32_t reg = *((volatile uint32_t*)DTM0);
    return reg >> 30 == 1;
}

void u3_write_field(uint32_t addr, uint32_t offset, uint32_t mask, uint32_t value) {
    uint32_t c = *((volatile uint32_t*)addr);
    uint32_t n = (c & (~mask)) | ((value << offset) & mask);
    __asm__ __volatile__("dsb\nisb":::"memory");
    *((volatile uint32_t*)addr) = n;
    __asm__ __volatile__("dsb\nisb":::"memory");
}

uint32_t mmio_read32(uint32_t addr) {
    __asm__ __volatile__("dsb":::"memory");
    uint32_t x = *((volatile uint32_t*)addr);
    return x;
}

void mmio_write32(uint32_t addr, uint32_t value) {
    *((volatile uint32_t*)addr) = value;
    __asm__ __volatile__("dsb\nisb":::"memory");
}

void test_uart() {
    volatile uint32_t* base = (volatile uint32_t*)0x11002000; // apuart0
    //volatile uint32_t* base = (volatile uint32_t*)0x11003000; // apuart1
    volatile uint32_t* lsr = (volatile uint32_t*)((uint32_t)base + 0x14);
#define THRE (1 << 5)

    video_printf("Beginning uart test");
    int i = 0;
    while(1) {
        *((volatile uint32_t*)0x110026E0) = *((volatile uint32_t*)0x110026E0) | 0x80000;
        uint32_t x = *lsr;
        video_printf("STILL RUNNING %d LSR=0x%x\n",i,x);

        while((x & THRE) == 0){}
        *base = 'H';

        i++;
    }
}

#if 0
bool is_uart0_on(bool print) {
    uint32_t status = *((volatile uint32_t*)(INFRASYS_BASE + 0x90));
    video_printf("status = 0x%x\n",status & 0x400000);
    bool is_on = (status & 0x400000) == 0;
    if (print)
        video_printf("UART0 state: %s\n", is_on ? "ON" : "OFF");
    return is_on;
}

void poweroff_uart0() {
    *((volatile uint32_t*)(INFRASYS_BASE + 0x80)) = 0x400000;
    __asm__ __volatile__("dsb":::"memory");
}

void test_uart_on_off() {
    is_uart0_on(true);
    poweroff_uart0();
    is_uart0_on(true);

    mdelay(500);
    *((volatile uint32_t*)0x110026E0) = *((volatile uint32_t*)0x110026E0) | 0x80000;
    test_uart();
    video_printf("HANG\n");
    for(;;);
}
#endif

bool is_ssusb_top_on(bool print) {
    uint32_t status = *((volatile uint32_t*)(INFRASYS_BASE + 0xac));
    video_printf("status = 0x%x\n",status & 2);

    bool is_on = (status & 2) == 0;
    if (print)
        video_printf("SSUSB_TOP status: %s\n", is_on ? "ON" : "OFF");
    return is_on;
}

void ssusb_top_off() {
    *((volatile uint32_t*)(INFRASYS_BASE + 0xa4)) = 2;
}

void enable_uart() {
    if (is_usb_uart_enabled()) {
        video_printf("UART already enabled");
        return;
    }

    /*mmio_write32(ACR6, mmio_read32(ACR6) & ~RG_USB20_BC11_SW_EN);
    mmio_write32(DTM0, mmio_read32(DTM0) | FORCE_UART_EN | FORCE_UART_TX_OE | FORCE_UART_BIAS_EN | (1 << RG_UART_MODE_OFST));
    mmio_write32(DTM1, mmio_read32(DTM1) | RG_UART_EN | RG_UART_TX_OE | RG_UART_BIAS_EN);
    mmio_write32(ACR4, mmio_read32(ACR4) | RG_USB20_DM_100K_EN);
    mmio_write32(DTM0, mmio_read32(DTM0) | RG_SUSPENDM);
    mmio_write32(DTM0, mmio_read32(DTM0) | FORCE_SUSPENDM);*/

#if 0
    U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_USBPHYACR6, RG_USB20_BC11_SW_EN_OFST, RG_USB20_BC11_SW_EN, 0);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM0, RG_SUSPENDM_OFST, RG_SUSPENDM, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM0, FORCE_SUSPENDM_OFST, FORCE_SUSPENDM, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM0, RG_UART_MODE_OFST, RG_UART_MODE, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM1, RG_UART_EN_OFST, RG_UART_EN, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM0, FORCE_UART_EN_OFST, FORCE_UART_EN, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM0, FORCE_UART_TX_OE_OFST, FORCE_UART_TX_OE, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM0, FORCE_UART_BIAS_EN_OFST, FORCE_UART_BIAS_EN, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM1, RG_UART_EN_OFST, RG_UART_EN, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM1, RG_UART_TX_OE_OFST, RG_UART_TX_OE, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYDTM1, RG_UART_BIAS_EN_OFST, RG_UART_BIAS_EN, 1);
	U3PhyWriteField32((phys_addr_t) (uintptr_t) U3D_U2PHYACR4, RG_USB20_DM_100K_EN_OFST, RG_USB20_DM_100K_EN, 1);
#endif

    u3_write_field(ACR6, RG_USB20_BC11_SW_EN_OFST, RG_USB20_BC11_SW_EN, 0);
    u3_write_field(DTM0, RG_SUSPENDM_OFST, RG_SUSPENDM, 1);
    u3_write_field(DTM0, FORCE_SUSPENDM_OFST, FORCE_SUSPENDM, 1);
    u3_write_field(DTM0, RG_UART_MODE_OFST, RG_UART_MODE, 1);
    u3_write_field(DTM1, RG_UART_EN_OFST, RG_UART_EN, 1);
	u3_write_field(DTM0, FORCE_UART_EN_OFST, FORCE_UART_EN, 1);
	u3_write_field(DTM0, FORCE_UART_TX_OE_OFST, FORCE_UART_TX_OE, 1);
	u3_write_field(DTM0, FORCE_UART_BIAS_EN_OFST, FORCE_UART_BIAS_EN, 1);
    u3_write_field(DTM1, RG_UART_TX_OE_OFST, RG_UART_TX_OE, 1);
	u3_write_field(DTM1, RG_UART_BIAS_EN_OFST, RG_UART_BIAS_EN, 1);

	u3_write_field(ACR4, RG_USB20_DM_100K_EN_OFST, RG_USB20_DM_100K_EN, 1);
    
    video_printf("acr4 = 0x%08x\n", mmio_read32(ACR4));
    is_ssusb_top_on(true);
    ssusb_top_off();
    is_ssusb_top_on(true);
    mdelay(50);

    while(1) {
        video_printf("acr4 = 0x%08x\n", mmio_read32(ACR4));
    }

    *((volatile uint32_t*)0x110026E0) = *((volatile uint32_t*)0x110026E0) | 0x80000;

    //test_uart();
    for(;;);

#if 0
    uint32_t active_uart = get_active_uart();
    video_printf("active_uart = 0x%x\n", active_uart);

    uint32_t x = *((volatile uint32_t*)0x10001084);
    video_printf("0x%x 0x%x\n", x, x & 0x400000);
    //test();
    *((volatile uint32_t*)0x10001084) = x | 0x400000u;
    __asm__ __volatile__("dsb":::"memory");

    x = *((volatile uint32_t*)0x10001084);
    video_printf("0x%x 0x%x\n", x, x & 0x400000);


    *((volatile uint32_t*)0x110026E0) = *((volatile uint32_t*)0x110026E0) | 0x80000;
#endif
    //for(;;);

    //mmio_read32(0x11002000 + 0xB0);
    //mmio_write32(0x11002000 + 0xB0,1);

    //if (is_usb_uart_enabled())
    //    video_printf("UART enabling OK");
    //else video_printf("UART enabling ERROR");

    //test_uart();
    //video_printf("HANG");
    //for(;;);
}
#endif

extern bool in_uart_mode;
int32_t phy_init_soc();
int32_t u2_slew_rate_calibration();
void usb20_pll_settings(bool host, bool forceOn);
void disable_interrupts();
void _ex_mu3d_hal_ssusb_en(void);
void mu3d_hal_rst_dev(void);

void usb_phy_switch_to_uart(void);
void usb_phy_switch_to_usb(void);
void uart_usb_switch_dump_register(void);
void usb_phy_recover(unsigned int clk_on);

int pmic_set_register_value() {
    return 0;
}

int pwrap_hal_init() {
    return 0;
}

void main(void)
{
    video_printf("hello from linuxboot\n");
    dprintf("hello from linuxboot\n");

    #if 0
    disable_interrupts();
    //linuxboot_pmic_init();
    video_printf("Enabling UART ... ");
    //usb_phy_switch_to_uart();
    //usb_phy_switch_to_usb();
    //usb_phy_switch_to_uart();
    //usb_phy_recover(false);

    in_uart_mode = true;
    phy_init_soc();
    usb_phy_recover(true);
    u2_slew_rate_calibration();
    _ex_mu3d_hal_ssusb_en();
    usb20_pll_settings(false, false);
    mu3d_hal_rst_dev();
    
    video_printf("DONE\n");

    //uart_usb_switch_dump_register();

    //video_printf("HANG\n");
    //for(;;);
    #endif
    // enable printk to uart
    *((volatile uint8_t*)(0x46074228+20)) = '0';

    part_dev_t *dev = mt_part_get_device();
    read_org = dev->read;
    dev->read = read;

    part_t *boot = mt_part_get_partition("boot");
    if (boot) {
        g_boot = boot->start_sect;
    } else {
        video_printf("boot partition not found\n");
        g_boot = 0xFFFFFFFF;
    }

    part_t *recovery = mt_part_get_partition("recovery");
    if (recovery) {
        g_recovery = boot->start_sect;
    } else {
        video_printf("recovery partition not found\n");
        g_recovery = 0xFFFFFFFF;
    }

    uint32_t *g_boot_mode = (uint32_t *)0x4609D120;
    video_printf("boot mode = %d\n", *g_boot_mode);

    //((uint32_t*)0x4609D120)[0] = 2; // boot recovery

    //trap(0x46027BAC);
    //trap(0x46027C98);
    //trap(0x46027D9E);
    //trap(0x46018A84);

    //trap(0x46027ED2);

    /*((uint32_t*)0x44000000)[0] = 0xAABBAABB;
    ((uint32_t*)0x44000000)[1] = 0xBBCCBBCC;
*/
    //*((uint32_t*)0x46018A84) = 0x4685;
    //*((uint32_t*)0x46018A86) = 0xffff;

    //trap(0x460297F6);

    // patches to aid in debugging

    // trap at infinite loop after "decompress kernel image fail ..."
    *((uint16_t*)0x46027EFC) = 0xFFFF;
    // trap at infinite loop after "can't find device tree ..."
    *((uint16_t*)0x46027ED2) = 0xFFFF;
    // trap at infinite loop after "root node search failed ..."
    *((uint16_t*)0x460290A0) = 0xFFFF;
    // trap at infinite loop after "add subnode memory failed ..."
    *((uint16_t*)0x46029096) = 0xFFFF;
    // trap at infinite loop after "/memory node does not exist ..."
    *((uint16_t*)0x46029030) = 0xFFFF;
    // trap at infinite loop after "64 bit kernel can't boot at ..."
    *((uint16_t*)0x46027F28) = 0xFFFF;


    // trap after "booting linux @ %p ..." (for debugging)
    // *((uint16_t*)0x46028720) = 0xFFFF;

    secure_boot_disable();

    uint32_t *g_boot_hdr = (uint32_t *)0x46120A9C;
    *g_boot_hdr = 0;

    video_printf("Loading DTB ... ");
    switch (*g_boot_mode)
    {
    case 0:
        bldr_load_dtb("boot");
        break;
    case 2:
        bldr_load_dtb("recovery");
        break;

    default:
        video_printf("Unknown boot mode (%d): doing normal boot", *g_boot_mode);
        *g_boot_mode = 0;
        bldr_load_dtb("boot");
        break;
    }
    video_printf("DONE\n");

//    enable_uart();

    mt_boot_init();
    for (;;)
        ;
}
